﻿
using System;
using System.Collections;
using System.Collections.Generic;
using System.Linq;
using System.Text;

using System.Diagnostics.Contracts;
using System.Reflection;
using HeapModel.ConcreteMemory;
using HeapModel.RegionPartition;
using HeapModel.ReducedGraph;

namespace HeapProfiler.Runtime
{
	public static class HeapWalk
	{
		/// <summary>
		/// Control the output of the profiler information, defaults to system out.
		/// </summary>
		public static System.IO.TextWriter PROFILER_OUTPUT;

		public static Random s_sameplev;
		public static System.Threading.Mutex s_mtx;
		public static bool s_activesample;

		public static string DumpDirectory;
		public static System.Diagnostics.Stopwatch TimerBuild;
		public static System.Diagnostics.Stopwatch TimerAbstract;

		static HeapWalk()
		{
			s_sameplev = new Random(1);
			s_mtx = new System.Threading.Mutex();
			s_activesample = false;

			TimerBuild = new System.Diagnostics.Stopwatch();
			TimerAbstract = new System.Diagnostics.Stopwatch();

			paramarray = new Object[MAX_PARAM_COUNT];
			samplemask = new bool[MAX_PARAM_COUNT];
			names = new String[MAX_PARAM_COUNT];

			PROFILER_OUTPUT = Console.Out;
			AppDomain.CurrentDomain.ProcessExit += new EventHandler(OnProcessExit);

			DumpDirectory = System.Environment.CurrentDirectory;
		}

		public static void OnProcessExit(object sender, EventArgs e)
		{
			if(DumpDirectory == null || (!System.IO.Directory.Exists(DumpDirectory)))
				return;

			DumpHeapProfile(DumpDirectory);
		}

		/////////////////////////////////////////////////////////////////////
		//Support for the heap walks done in the CmdLineExtractor bits.
		//Code to programatically set the roots we want to dump.

		////
		//Load up and manage the roots that we are going to do the heap dump from.

		/// <summary>
		/// The maximum number of roots allowed.
		/// </summary>
		private const int MAX_PARAM_COUNT = 12;

		/// <summary>
		/// The number of roots that we want to dump on the next dump call (see <c>InitDumpState</c>)
		/// </summary>
		private static int paramc;

		/// <summary>
		/// The references to the objects we want to dump (we will start walking the heap via reflection on these objects);
		/// </summary>
		private static Object[] paramarray;

		/// <summary>
		/// The true if we are going to walk from the location, false otherwise.
		/// </summary>
		private static bool[] samplemask;

		/// <summary>
		/// The names YOU want to give these roots.
		/// </summary>
		private static String[] names;

		/// <summary>
		/// Indicate how many roots will be initialized for the next heap walk.
		/// </summary>
		public static void InitDumpState(int argc)
		{
			Contract.Requires(0 <= argc);
			Contract.Requires(argc < MAX_PARAM_COUNT);

			if(s_activesample)
				return;

			s_mtx.WaitOne();

			paramc = argc;
			Array.Clear(paramarray, 0, MAX_PARAM_COUNT);
			Array.Clear(samplemask, 0, MAX_PARAM_COUNT);
			Array.Clear(names, 0, MAX_PARAM_COUNT);
		}

		/// <summary>
		/// Clear all the root info.
		/// </summary>
		public static void ClearDumpState()
		{
			if(s_activesample)
				return;

			Array.Clear(paramarray, 0, MAX_PARAM_COUNT);
			Array.Clear(samplemask, 0, MAX_PARAM_COUNT);
			Array.Clear(names, 0, MAX_PARAM_COUNT);

			s_mtx.ReleaseMutex();
		}

		/// <summary>
		/// Set the ith root parameter for the next heap walk.
		/// </summary>
		/// <param name="pos">The position in the root set to put the info.</param>
		/// <param name="oval">The actual object to start at (or null if the object is null or something we are not sampling from).</param>
		/// <param name="useinsample">True if we should walk from this argument.</param>
		/// <param name="pname">The name we want to give the root.</param>
		public static void SetIthParamInfo(int pos, Object oval, bool useinsample, String pname)
		{
			Contract.Requires(0 <= pos);
			Contract.Requires(!String.IsNullOrEmpty(pname));

			if(s_activesample)
				return;

			paramarray[pos] = oval;
			samplemask[pos] = useinsample;
			names[pos] = pname;
		}

		/////////////////////////////////////////////////////////////////////
		//Support for the heap walks done in the CmdLineProfiler bits.

		private static Dictionary<int, Tuple<int, string, ReducedHeapGraph>> s_profilerentries = new Dictionary<int, Tuple<int, string, ReducedHeapGraph>>();

		private static Dictionary<int, double> s_methodsampledict = new Dictionary<int, double>();

		/// <summary>
		/// The random sample rate.
		/// </summary>
		private static double s_defaultsamplerate = 0.1;

		/// <summary>
		/// The threshold we use to decide if the new heap is sufficiently larger than the heaps we have seen.
		/// </summary>
		private static Dictionary<int, int> s_nextsamplesizedict = new Dictionary<int, int>();

		private static void InitializeShouldSample(int token)
		{ s_nextsamplesizedict.Add(token, 0); }

		private static int GetShouldSampleSize(int token)
		{ return s_nextsamplesizedict[token]; }

		private static void SetShouldSampleSize(int token, int newsize)
		{ s_nextsamplesizedict[token] = newsize; }

		private static ConcreteMemoryBuilder s_cmb = new ConcreteMemoryBuilder();

		/// <summary>
		/// Trigger a heap walk with abstraction and accumulation for performing memory profiling.
		/// You Must Set Up ArgList Info Before Calling this.
		/// </summary>
		public static void SampleMemProfileAtCallEntry(string fmname, int token)
		{
			if(s_activesample)
				return;

			if(!s_methodsampledict.ContainsKey(token))
			{
				s_methodsampledict.Add(token, s_defaultsamplerate);
				s_profilerentries.Add(token, new Tuple<int, string, ReducedHeapGraph>(-1, fmname, null));
				InitializeShouldSample(token);
			}

			double rs = s_sameplev.NextDouble();
			if(s_methodsampledict[token] >= rs)
			{
				s_activesample = true;

				long memuse = GC.GetTotalMemory(true);
				int nextsamplesize = GetShouldSampleSize(token);

				if(memuse < (0.9 * nextsamplesize))
					s_methodsampledict[token] = Math.Max(.00001, s_methodsampledict[token] / 10.0);
				else
				{
					//debug
					//Console.WriteLine("Sampling for size " + s_nextsamplesize + "...");

					//System.Diagnostics.Debugger.Launch();

					TimerBuild.Start();
					s_cmb.ResetConcreteHeap();
					ConcreteHeap theap = s_cmb.ExtractConcreteHeap(paramc, paramarray, samplemask, names);
					TimerBuild.Stop();

					int size = theap.AllObjects.Sum((ConcreteNode cobj) => (cobj.DataSizeInBytes + cobj.GlueSizeInBytes));
					if(size < nextsamplesize)
						s_methodsampledict[token] = Math.Max(.00001, s_methodsampledict[token] / 10.0);
					else
					{
						Console.WriteLine(String.Format("Sampling method {0} for size {1}...", fmname, nextsamplesize));

						TimerAbstract.Start();
						RegionGraph rpart = RegionGraph.Create(theap, s_cmb.AllTypesSet);
						Console.WriteLine("Done Region.");
						
						ReducedHeapGraph rdhg = ReducedHeapGraph.Create(rpart);
						Console.WriteLine("Done Reduced.");
						TimerAbstract.Stop();

						//System.Diagnostics.Debugger.Launch();
						//RegionGraph rg = RegionGraph.Create_NoNorm(theap, cmb.AllTypesSet);
						//rg.WriteDgml("tmp.dgml");

						Tuple<int, string, ReducedHeapGraph> profilerheap = new Tuple<int, string, ReducedHeapGraph>(size, fmname, rdhg);
						s_profilerentries[token] = profilerheap;

						HashSet<int> largeheaps = new HashSet<int>(s_profilerentries.Values.Select((Tuple<int, string, ReducedHeapGraph> rtup) => rtup.Item1));
						if(largeheaps.Count > 30)
						{
							List<int> lhl = new List<int>(largeheaps);
							lhl.Sort();
							lhl.Reverse();
							int mv = lhl[10];

							int[] keys = s_profilerentries.Keys.ToArray();
							for(int i = 0; i < keys.Length; ++i)
							{
								if(s_profilerentries[keys[i]].Item1 < mv)
								{
									Tuple<int, string, ReducedHeapGraph> ctt = s_profilerentries[keys[i]];
									s_profilerentries[keys[i]] = new Tuple<int, string, ReducedHeapGraph>(0, ctt.Item2, null);
								}
							}
						}

						int mnextsamplesize = (int)(size * 1.5);
						SetShouldSampleSize(token, mnextsamplesize);

						List<int> skeys = new List<int>(s_methodsampledict.Keys);
						foreach(int kk in skeys)
							s_methodsampledict[kk] = s_defaultsamplerate;

						Console.WriteLine("New Size " + mnextsamplesize + ".");
						Console.WriteLine("Done.");
					}

					//debug
					//Console.WriteLine("Done.");
				}

				s_activesample = false;
			}

			ClearDumpState();
		}

		/// <summary>
		/// Dump all of the accumulated heap state information into files in the given directory
		/// </summary>
		public static void DumpHeapProfile(String dumpdirectory)
		{
			if(PROFILER_OUTPUT != null)
				PROFILER_OUTPUT.WriteLine("Dumping Profiler Output as DGML...");

			System.Xml.XmlWriterSettings settings = new System.Xml.XmlWriterSettings();
			settings.Indent = true;
			settings.OmitXmlDeclaration = true;

			List<Tuple<int, string, ReducedHeapGraph>> bigheaps = new List<Tuple<int, string, ReducedHeapGraph>>(s_profilerentries.Values);
			bigheaps.RemoveAll((Tuple<int, string, ReducedHeapGraph> he) => he.Item1 < 64);
			bigheaps.Sort((Tuple<int, string, ReducedHeapGraph> he1, Tuple<int, string, ReducedHeapGraph> he2) => he2.Item1 - he1.Item1);

			string dumpdir2 = System.IO.Path.Combine(dumpdirectory, "HeapDumps\\");
			if(!System.IO.Directory.Exists(dumpdir2))
				System.IO.Directory.CreateDirectory(dumpdir2);

			for(int i = 0; i < 10 & i < bigheaps.Count; ++i)
			{
				Tuple<int, string, ReducedHeapGraph> ph = bigheaps[i];
				string hss;
				int maxsize = ph.Item1;
				if(maxsize < 1000)
					hss = maxsize + "B";
				else if(maxsize < 1000000)
					hss = (((double)maxsize) / 1000.0) + "KB";
				else
					hss = (((double)maxsize) / 1000000.0) + "MB";

				if(PROFILER_OUTPUT != null)
				{
					PROFILER_OUTPUT.WriteLine("Dumping Memory Profile Heap (of size " + hss + " ).");
					PROFILER_OUTPUT.WriteLine("Result in: " + ph.Item2 + ".");
				}

				List<char> allch = new List<char>(ph.Item2);
				HashSet<char> badc = new HashSet<char>(System.IO.Path.GetInvalidFileNameChars());
				allch.RemoveAll((char c) => badc.Contains(c));
				string rchn = new String(allch.ToArray());

				string dgmlpathb = System.IO.Path.Combine(dumpdir2, ph.Item2 + ".dgml");
				ph.Item3.WriteDgml(dgmlpathb, dgmlpathb.Replace(".dgml", "-base.dgml"));

				string xmlfilename = System.IO.Path.ChangeExtension(dgmlpathb, ".xml");
				using(var xw = System.Xml.XmlWriter.Create(xmlfilename, new System.Xml.XmlWriterSettings { Indent = true }))
					ph.Item3.BaseHG.CMem.XMLify(xw, new List<HeapModel.ProgInfo.TypeIdentifier>(s_cmb.AllTypesSet), new List<HeapModel.ProgInfo.OffsetIdentifier>(s_cmb.AllOffsetsSet));
				
				if(PROFILER_OUTPUT != null)
					PROFILER_OUTPUT.WriteLine("Done.");
			}

			if(PROFILER_OUTPUT != null)
			{
				PROFILER_OUTPUT.WriteLine("Total Extraction Time (seconds): " + TimerBuild.Elapsed.TotalSeconds);
				PROFILER_OUTPUT.WriteLine("Total Abstraction Time (seconds): " + TimerAbstract.Elapsed.TotalSeconds);
				PROFILER_OUTPUT.WriteLine();
			}
		}
	}
}
